import csv
import sys
import codecs

import matplotlib.pyplot as plt

from utils import *
from scipy.fftpack import fft,ifft
from sklearn.decomposition import PCA
from sklearn.neighbors import KNeighborsClassifier

trainData,trainLabels,validData,validLabels=getTrainValid()
testFiles,testData=getOriginalTest()

for i in range(len(trainData)):
    trainData[i]=abs(fft(trainData[i])[:int(trainDestLen/2+1)])
for i in range(len(validData)):
    validData[i]=abs(fft(validData[i])[:int(trainDestLen/2+1)])

pca=PCA(10)
pca.fit(trainData)

'''vars=pca.explained_variance_
plt.plot([i for i in range(len(vars))],vars)
plt.show()'''

trainData=pca.transform(trainData)
validData=pca.transform(validData)

knn=KNeighborsClassifier()
m3valid=codecs.open('./m3_valid.txt','w','UTF-8')

maxAcc=0
bestK=-1

accList=[]

for i in range(5,16):
    knn = KNeighborsClassifier(n_neighbors=i)
    knn.fit(trainData, trainLabels)
    predicted = knn.predict(validData)
    accuracy = (np.array(predicted) == validLabels).sum() * 1.0 / len(validLabels)
    accStr = str(i) + ":" + str(accuracy)
    print(accStr)
    m3valid.write(accStr + '\n')
    accList.append(accuracy)
    if accuracy > maxAcc:
        bestK = i
        maxAcc = accuracy

plt.plot([i for i in range(5,16)],accList)
plt.show()


m3valid.write(str(bestK)+':'+str(maxAcc))
m3valid.close()
print(bestK,maxAcc)

knn=KNeighborsClassifier(n_neighbors=bestK)
knn.fit(trainData,trainLabels)

testFiles,testData=getOriginalTest()
testPredicted=[]


for i in range(len(testData)):
    oneData=testData[i]
    dataLen=len(oneData)
    index=trainDestLen
    segments=[]
    while index<=dataLen:
        segments.append(abs(fft(oneData[index - trainDestLen:index])[:int(trainDestLen/2)+1]))
        index = index + trainDestLen
    segments=pca.transform(segments)
    onePreResult = knn.predict(segments)
    testPredicted.append(np.argmax(np.bincount(onePreResult)))

csvHeaders=['id','categories']
csvData=[[testFiles[i],testPredicted[i]] for i in range(len(testFiles))]

with open('method3.csv', 'w', encoding='utf-8', newline='') as f:
    writer = csv.writer(f)
    writer.writerow(csvHeaders)
    writer.writerows(csvData)